# -*- coding: UTF-8 -*-
"""
@项目名称：protein_add_ligand.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-19-10:33
"""

from Bio.PDB import PDBParser, PDBIO
from Bio.PDB.Structure import Structure
from Bio.PDB.Model import Model
from Bio.PDB.Chain import Chain
import os

class ProteinLigandComplexBuilder:
    """
    用于构建蛋白质-配体复合物的类
    """

    def __init__(self):
        """初始化PDB解析器"""
        self.parser = PDBParser(QUIET=True)  # 创建解析器并忽略警告信息

    def load_structure(self, file_path):
        """
        加载PDB文件

        参数:
            file_path: PDB文件路径

        返回:
            Bio.PDB.Structure对象或None(如果加载失败)
        """
        try:
            if not os.path.exists(file_path):
                print(f"错误: 文件 '{file_path}' 不存在")
                return None

            # 从文件名提取结构ID
            structure_id = os.path.basename(file_path).split('.')[0]

            # 解析PDB文件
            structure = self.parser.get_structure(structure_id, file_path)
            return structure

        except Exception as e:
            print(f"解析PDB文件时出错: {str(e)}")
            return None

    def add_ligand_to_protein(self, protein_structure, ligand_structure,
                              target_chain_id='A', ligand_residue_name='LIG'):
        """
        将配体添加到蛋白质结构中

        参数:
            protein_structure: 蛋白质的Structure对象
            ligand_structure: 配体的Structure对象
            target_chain_id: 目标链ID，配体将添加到此链
            ligand_residue_name: 配体的残基名称

        返回:
            添加了配体的新Structure对象
        """
        # 创建新的结构对象
        complex_structure = Structure("Protein_Ligand_Complex")

        # 复制蛋白质的所有模型
        for model in protein_structure:
            new_model = Model(model.get_id())
            # 复制蛋白质的所有链
            for chain in model:
                new_chain = Chain(chain.get_id())
                # 复制蛋白质的所有残基
                for residue in chain:
                    new_chain.add(residue.copy())
                # 如果是目标链，添加配体
                if chain.get_id() == target_chain_id:
                    # 获取配体的所有残基（通常配体只有一个残基）
                    for ligand_model in ligand_structure:
                        for ligand_chain in ligand_model:
                            for ligand_residue in ligand_chain:
                                # 创建配体残基的副本并修改残基名称
                                ligand_residue_copy = ligand_residue.copy()
                                ligand_residue_copy.resname = ligand_residue_name

                                # 为配体分配一个新的残基ID（通常使用一个较高的数字）
                                # 找到蛋白质中最大的残基ID
                                max_residue_id = 0
                                for res in new_chain:
                                    res_id = res.id[1]  # 残基ID是元组的第二个元素
                                    if res_id > max_residue_id:
                                        max_residue_id = res_id

                                # 设置配体的新残基ID
                                new_residue_id = (' ', max_residue_id + 1, ' ')  # (hetero flag, sequence number, insertion code)
                                ligand_residue_copy.id = new_residue_id

                                # 添加配体到链
                                new_chain.add(ligand_residue_copy)
                                print(f"已将配体添加到链 {target_chain_id}，残基ID: {new_residue_id}")
                new_model.add(new_chain)
            complex_structure.add(new_model)
        return complex_structure

    def save_complex(self, complex_structure, output_path):
        """
        保存复合物结构到PDB文件
        参数:
            complex_structure: 复合物Structure对象
            output_path: 输出文件路径
        """
        io = PDBIO()
        io.set_structure(complex_structure)
        io.save(output_path)
        print(f"复合物结构已保存到: {output_path}")

# 示例用法
def main():
    # 创建复合物构建器实例
    builder = ProteinLigandComplexBuilder()

    # 指定蛋白质和配体PDB文件路径
    protein_pdb_file = "/4T/AIDDserver/DataSrc/lpa_test/4bvw/extracted_chains_class/chain_蛋白质_A.pdb"  # 请替换为你的蛋白质PDB文件路径
    ligand_pdb_file = "/4T/AIDDserver/DataSrc/lpa_test/4bvw/extracted_chains_class/chain_配体_D.pdb"   # 请替换为你的配体PDB文件路径
    # 保存复合物结构
    output_file = "/4T/AIDDserver/DataSrc/lpa_test/4bvw/protein_ligand_complex.pdb"
    # 加载蛋白质和配体结构
    protein_structure = builder.load_structure(protein_pdb_file)
    if protein_structure is None:
        print("无法加载蛋白质结构")
        return
    ligand_structure = builder.load_structure(ligand_pdb_file)
    if ligand_structure is None:
        print("无法加载配体结构")
        return
    # 将配体添加到蛋白质中
    complex_structure = builder.add_ligand_to_protein(
        protein_structure,
        ligand_structure,
        target_chain_id='A',  # 配体将添加到A链
        ligand_residue_name='LIG'  # 配体残基名称
    )
    builder.save_complex(complex_structure, output_file)
    print("操作完成!")

if __name__ == "__main__":
    main()